import torch
import torch.nn as nn
import torch.nn.functional as F


class ResBlock(nn.Module):
    def __init__(self, in_channel, out_channel, stride=1):
        super(ResBlock, self).__init__()
        # 基本算子
        # 主干分支
        self.layer = nn.Sequential(
            # 只做一次stride下采样
            nn.Conv2d(in_channel, out_channel, kernel_size=3, stride=stride, padding=1),
            nn.BatchNorm2d(out_channel),
            nn.ReLU(),

            nn.Conv2d(out_channel, out_channel, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(out_channel),
            nn.ReLU()
        )

        # 如果通道数发生了改变 或者 步长大于1 说明主干分支对输入的x做了改变 需要增加跳链分支
        # 跳链分支
        self.shortcut = nn.Sequential()
        if in_channel != out_channel or stride > 1:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channel, out_channel, kernel_size=3, stride=stride, padding=1),
                nn.BatchNorm2d(out_channel)
            )

    def forward(self, x):
        out_1 = self.layer(x)
        out_2 = self.shortcut(x)
        out = out_1 + out_2
        out = F.relu(out)
        return out


class ResNet(nn.Module):

    # 循环添加resBlock层
    def make_layer(self, block, out_channel, stride, num_block):
        layer_list = []
        for i in range(num_block):
            if i == 0:
                in_stride = stride
            else:
                in_stride = 1
            layer_list.append(block(self.in_channel, out_channel, in_stride))
            # 结束后修改self_in_channel
            self.in_channel = out_channel

        return nn.Sequential(*layer_list)

    def __init__(self):
        super(ResNet, self).__init__()

        # =====================================================================
        # 卷积层 3*28*28 -》 64*14*14
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU()
        )

        # 一个resBlock
        # self.layer1_1 = ResBlock(in_channel=32, out_channel=64, stride=2)
        # self.layer1_2 = ResBlock(in_channel=64, out_channel=64, stride=1)
        # self.layer1_3 = ResBlock(in_channel=64, out_channel=64, stride=1)
        # 一个resBlock
        # self.layer2_1 = ResBlock(in_channel=64, out_channel=128, stride=2)
        # self.layer2_2 = ResBlock(in_channel=128, out_channel=128, stride=1)
        # self.layer2_3 = ResBlock(in_channel=128, out_channel=128, stride=1)
        # 自定义方法循环resBlock层
        self.in_channel = 32
        self.layer1 = self.make_layer(block=ResBlock, out_channel=64, stride=2, num_block=2)
        self.layer2 = self.make_layer(block=ResBlock, out_channel=128, stride=2, num_block=2)
        self.layer3 = self.make_layer(block=ResBlock, out_channel=256, stride=2, num_block=2)
        self.layer4 = self.make_layer(block=ResBlock, out_channel=512, stride=2, num_block=2)
        self.fc = nn.Linear(512, 10)

    def forward(self, x):
        out = self.conv1(x)
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        # 转化为batch-size*512*1*1
        out = F.avg_pool2d(out, 2)
        # 转为为batch-size*(512*1*1)
        out = out.view(out.size(0), -1)
        out = self.fc(out)
        return out


def resNet():
    return ResNet()
